#!/usr/bin/env python3
import zipfile
import xml.etree.ElementTree as ET
from collections import Counter
import sys

DOCX_PATH = 'final-report.docx'
SEARCH_STYLE = 'h3-song-bold'

def extract_xml(docx_path, xml_path):
    with zipfile.ZipFile(docx_path, 'r') as docx:
        with docx.open(xml_path) as xml_file:
            return xml_file.read()

def count_styles_in_document(xml_bytes):
    ns = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'}
    root = ET.fromstring(xml_bytes)
    style_ids = []

    # 搜索所有段落 <w:p> 中的 <w:pPr>/<w:pStyle w:val="..."/>
    for p in root.findall('.//w:p', ns):
        pPr = p.find('w:pPr', ns)
        if pPr is not None:
            pStyle = pPr.find('w:pStyle', ns)
            if pStyle is not None:
                style_val = pStyle.attrib.get(f'{{{ns["w"]}}}val')
                if style_val:
                    style_ids.append(style_val)
    return Counter(style_ids)

def find_style_definition(xml_bytes, style_id):
    ns = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'}
    root = ET.fromstring(xml_bytes)
    # 找 <w:style w:styleId="style_id">
    for style in root.findall('.//w:style', ns):
        sid = style.attrib.get(f'{{{ns["w"]}}}styleId')
        if sid == style_id:
            # 找到后返回该节点字符串或者简要信息
            return ET.tostring(style, encoding='unicode')
    return None

def main():
    try:
        doc_xml = extract_xml(DOCX_PATH, 'word/document.xml')
        styles_xml = extract_xml(DOCX_PATH, 'word/styles.xml')
    except KeyError as e:
        print(f"Error: {e} - Ensure the .docx file contains the expected XML parts.")
        sys.exit(1)

    print(f"统计正文中所有段落样式及出现次数：")
    counter = count_styles_in_document(doc_xml)
    for style_id, count in counter.most_common():
        print(f"{style_id}: {count}")

    print("\n查找样式定义：", SEARCH_STYLE)
    definition = find_style_definition(styles_xml, SEARCH_STYLE)
    if definition:
        print(definition)
    else:
        print(f"样式 '{SEARCH_STYLE}' 未找到定义。")

if __name__ == "__main__":
    main()
